import torch
from torch import nn

class CnnLayer(nn.Module):
    def __init__(self,c_in,c_out):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(c_in,c_out,3,1,1),
            nn.BatchNorm2d(c_out),
            nn.LeakyReLU(),
            nn.Conv2d(c_out,c_out,3,1,1),
            nn.BatchNorm2d(c_out),
            nn.LeakyReLU()
        )
    def forward(self,x):
        return self.layer(x)
class DownSampling(nn.Module):
    def __init__(self,c):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(c,c,3,2,1),
            nn.LeakyReLU()
        )
    def forward(self,x):
        return self.layer(x)

class UpSampling(nn.Module):
    def __init__(self,c):
        super().__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(c,c//2,1,1),
            nn.LeakyReLU()
        )
    def forward(self,x,out):
        up = nn.functional.interpolate(x,x.shape[-1]*2)
        x = self.layer(up)
        return torch.cat((x,out),1)

class MainNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.c1 = CnnLayer(3,64)
        self.d1 = DownSampling(64)
        self.c2 = CnnLayer(64,128)
        self.d2 = DownSampling(128)
        self.c3 = CnnLayer(128,256)
        self.d3 = DownSampling(256)
        self.c4 = CnnLayer(256,512)
        self.d4 = DownSampling(512)
        self.c5 = CnnLayer(512,1024)
        self.u1 = UpSampling(1024)
        self.c6 = CnnLayer(1024,512)
        self.u2 = UpSampling(512)
        self.c7 = CnnLayer(512,256)
        self.u3 = UpSampling(256)
        self.c8 = CnnLayer(256,128)
        self.u4 = UpSampling(128)
        self.c9 = CnnLayer(128,64)
        self.sigmoid = nn.Sigmoid()
        self.pre = nn.Conv2d(64,22,3,1,1)
    def forward(self,x):
        R1 = self.c1(x)
        R2 = self.c2(self.d1(R1))
        R3 = self.c3(self.d2(R2))
        R4 = self.c4(self.d3(R3))
        Y1 = self.c5(self.d4(R4))

        O1 = self.c6(self.u1(Y1,R4))
        O2 = self.c7(self.u2(O1, R3))
        O3 = self.c8(self.u3(O2, R2))
        O4 = self.c9(self.u4(O3, R1))
        return self.sigmoid(self.pre(O4))
if __name__ == '__main__':

    a = torch.randn(2,3,256,256)
    net = MainNet()
    print(net(a).shape)
